import Foundation
import Combine
import SwiftUI

// 聊天消息模型
struct ChatMessage: Identifiable, Codable, Equatable {
    var id = UUID()
    var content: String
    var role: MessageRole
    var timestamp: Date
    var isLoading: Bool = false
    
    enum MessageRole: String, Codable, Equatable {
        case user
        case assistant
        case system
    }
    
    // 实现Equatable协议
    static func == (lhs: ChatMessage, rhs: ChatMessage) -> Bool {
        return lhs.id == rhs.id &&
               lhs.content == rhs.content &&
               lhs.role == rhs.role &&
               lhs.timestamp == rhs.timestamp &&
               lhs.isLoading == rhs.isLoading
    }
}

// 聊天会话模型
struct ChatSession: Identifiable, Codable, Equatable {
    var id = UUID()
    var title: String
    var messages: [ChatMessage]
    var createdAt: Date
    var updatedAt: Date
    
    // 实现Equatable协议
    static func == (lhs: ChatSession, rhs: ChatSession) -> Bool {
        return lhs.id == rhs.id &&
               lhs.title == rhs.title &&
               lhs.messages == rhs.messages &&
               lhs.createdAt == rhs.createdAt &&
               lhs.updatedAt == rhs.updatedAt
    }
}

class ChatViewModel: ObservableObject {
    @Published var sessions: [ChatSession] = []
    @Published var currentSession: ChatSession
    @Published var isLoading: Bool = false
    @Published var errorMessage: String?
    @Published var currentModelType: AIModelType = .gpt4o // 默认使用GPT-4o
    
    private var cancellables = Set<AnyCancellable>()
    
    init() {
        // 初始化一个默认会话
        currentSession = ChatSession(
            title: "新对话",
            messages: [],
            createdAt: Date(),
            updatedAt: Date()
        )
        
        loadSessions()
        
        // 如果没有会话，创建一个默认会话
        if sessions.isEmpty {
            createNewSession()
        } else {
            // 使用最近的会话作为当前会话
            currentSession = sessions.first!
        }
    }
    
    // 创建新会话
    func createNewSession() {
        let newSession = ChatSession(
            title: "新对话 \(sessions.count + 1)",
            messages: [],
            createdAt: Date(),
            updatedAt: Date()
        )
        
        sessions.insert(newSession, at: 0)
        currentSession = newSession
        
        saveSessions()
    }
    
    // 切换会话
    func switchSession(to sessionId: UUID) {
        if let session = sessions.first(where: { $0.id == sessionId }) {
            currentSession = session
        }
    }
    
    // 重命名会话
    func renameSession(id: UUID, newTitle: String) {
        if let index = sessions.firstIndex(where: { $0.id == id }) {
            sessions[index].title = newTitle
            
            if currentSession.id == id {
                currentSession.title = newTitle
            }
            
            saveSessions()
        }
    }
    
    // 删除会话
    func deleteSession(id: UUID) {
        sessions.removeAll(where: { $0.id == id })
        
        if currentSession.id == id {
            if let firstSession = sessions.first {
                currentSession = firstSession
            } else {
                createNewSession()
            }
        }
        
        saveSessions()
    }
    
    // 清除当前会话
    func clearCurrentSession() {
        currentSession.messages.removeAll()
        saveSessions()
    }
    
    // 发送消息
    func sendMessage(_ content: String) {
        guard !content.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
        
        // 添加用户消息
        let userMessage = ChatMessage(
            content: content,
            role: .user,
            timestamp: Date()
        )
        
        currentSession.messages.append(userMessage)
        
        // 添加一个加载中的助手消息
        let loadingMessage = ChatMessage(
            content: "",
            role: .assistant,
            timestamp: Date(),
            isLoading: true
        )
        
        let loadingIndex = currentSession.messages.count
        currentSession.messages.append(loadingMessage)
        
        // 更新会话
        currentSession.updatedAt = Date()
        if currentSession.messages.count == 2 {
            // 如果是第一条消息，更新会话标题
            currentSession.title = content.prefix(20).trimmingCharacters(in: .whitespacesAndNewlines)
        }
        
        saveSessions()
        
        // 获取核心记忆内容
        var systemPrompt = ""
        let coreMemories = MemoryService.shared.getCoreMemories()
        
        // 添加人类记忆
        if let humanMemory = coreMemories.first(where: { $0.category == .human }) {
            systemPrompt += "关于用户的信息：\n\(humanMemory.content)\n\n"
        }
        
        // 添加角色记忆
        if let personaMemory = coreMemories.first(where: { $0.category == .persona }) {
            systemPrompt += "你的角色设定：\n\(personaMemory.content)\n\n"
        }
        
        // 设置加载状态
        isLoading = true
        errorMessage = nil
        
        // 获取除了最新加载消息之外的所有消息
        let messagesToSend = Array(currentSession.messages.dropLast())
        
        // 使用OpenAIService发送消息，使用真正的流式输出
        OpenAIService.shared.sendRealTimeStreamMessage(
            messages: messagesToSend,
            systemPrompt: systemPrompt,
            modelType: currentModelType,
            onReceive: { [weak self] chunk in
                guard let self = self else { return }
                
                // 更新助手消息内容
                if self.currentSession.messages.indices.contains(loadingIndex) {
                    self.currentSession.messages[loadingIndex].content += chunk
                    self.currentSession.messages[loadingIndex].isLoading = false
                    self.objectWillChange.send() // 强制UI更新
                }
            },
            onComplete: { [weak self] error in
                guard let self = self else { return }
                
                self.isLoading = false
                
                if let error = error {
                    self.errorMessage = "发送消息失败: \(error.localizedDescription)"
                    
                    // 移除加载中的消息
                    if self.currentSession.messages.indices.contains(loadingIndex) {
                        self.currentSession.messages.remove(at: loadingIndex)
                    }
                } else {
                    // 完成后更新长期记忆
                    MemoryService.shared.manageLongTermMemory(from: self.currentSession.messages)
                }
                
                self.saveSessions()
                self.objectWillChange.send() // 强制UI更新
            }
        )
    }
    
    // 重新生成最后一个响应
    func regenerateLastResponse() {
        guard !currentSession.messages.isEmpty else { return }
        
        // 如果最后一条消息是助手消息，则移除它
        if currentSession.messages.last?.role == .assistant {
            currentSession.messages.removeLast()
        }
        
        // 如果还有消息，且最后一条是用户消息，则重新发送
        if !currentSession.messages.isEmpty, let lastUserMessage = currentSession.messages.last, lastUserMessage.role == .user {
            sendMessage(lastUserMessage.content)
        }
    }
    
    // 创建归档记忆
    func createArchivalMemory(from content: String, category: Memory.MemoryCategory) {
        let memory = Memory(
            content: content,
            type: .archival,
            category: category,
            createdAt: Date(),
            updatedAt: Date()
        )
        
        MemoryService.shared.saveArchivalMemory(memory)
    }
    
    // 更新核心记忆
    func updateCoreMemory(id: UUID, content: String) {
        if let memory = MemoryService.shared.getCoreMemories().first(where: { $0.id == id }) {
            let updatedMemory = Memory(
                id: memory.id,
                content: content,
                type: memory.type,
                category: memory.category,
                createdAt: memory.createdAt,
                updatedAt: Date()
            )
            
            MemoryService.shared.saveCoreMemory(updatedMemory)
        }
    }
    
    // 保存会话
    private func saveSessions() {
        // 更新当前会话在会话列表中的状态
        if let index = sessions.firstIndex(where: { $0.id == currentSession.id }) {
            sessions[index] = currentSession
        } else {
            sessions.append(currentSession)
        }
        
        // 保存到UserDefaults
        if let encoded = try? JSONEncoder().encode(sessions) {
            UserDefaults.standard.set(encoded, forKey: "ChatSessions")
        }
    }
    
    // 加载会话
    private func loadSessions() {
        if let savedSessions = UserDefaults.standard.data(forKey: "ChatSessions"),
           let decoded = try? JSONDecoder().decode([ChatSession].self, from: savedSessions) {
            sessions = decoded
        }
    }
    
    // 设置当前使用的模型
    func setCurrentModel(_ model: AIModelType) {
        self.currentModelType = model
    }
    
    // 获取当前模型配置
    func getCurrentModelConfig() -> AIModelConfig {
        var config = AIModelConfig()
        config.modelType = currentModelType
        config.provider = currentModelType.provider
        return config
    }
} 